{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "from IPython.display import Image\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from pyodesys.symbolic import ScaledSys\n",
    "from chempy import ReactionSystem\n",
    "from chempy.kinetics.ode import get_odesys\n",
    "from chempy.kinetics.analysis import plot_reaction_contributions, dominant_reactions_graph\n",
    "from chempy.kinetics._native import get_native\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['OMP_NUM_THREADS'] = '1'\n",
    "print(os.environ['OMP_NUM_THREADS'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rate constants chosen rather arbitrarily for demonstration purposes\n",
    "system_str = \"\"\"\n",
    "  2 Fe+2 +  H2O2 -> 2 Fe+3     + 2 OH-        ; 42\n",
    "  2 Fe+3 +  H2O2 -> 2 Fe+2     +   O2  + 2 H+ ; 17\n",
    "      H+ +   OH- ->   H2O                     ; %(kprot_H2O)s\n",
    "             H2O ->   H+       +   OH-        ; %(kprot_H2O)s * %(Kw)s\n",
    "   Fe+3  + 2 H2O ->   FeOOH(s) + 3 H+         ; 1\n",
    "FeOOH(s) + 3  H+ ->   Fe+3     + 2 H2O        ; 2.5\n",
    "            H2O2 ->   H+       + HO2-         ; %(kprot_H2O2)s * %(Ka_H2O2)s\n",
    "      H+ +  HO2- ->   H2O2                    ; %(kprot_H2O2)s\n",
    "    H2O2 +   OH- ->   HO2-     + H2O          ; %(kdepr_H2O2)s\n",
    "    HO2- +   H2O ->   H2O2     + OH-          ; %(kdepr_H2O2)s / %(Ka_H2O2)s * %(Kw)s\n",
    "\"\"\" % dict(Kw='1e-14', kprot_H2O='1e10', Ka_H2O2='10**-11.65', kprot_H2O2='5e10', kdepr_H2O2='1.3e10')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rsys = ReactionSystem.from_string(system_str, name='Iron(II)/Iron(III) speciation')  # \"[H2O]\" = 1.0 (actually 55.4 at RT)\n",
    "odesys, extra = get_odesys(rsys, SymbolicSys=ScaledSys, dep_scaling=1e6)\n",
    "rsys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tout = sorted(np.concatenate((np.linspace(0, 1003), np.logspace(-8, 3))))\n",
    "c0 = defaultdict(float, {'Fe+2': 0.05, 'H2O2': 0.1, 'H2O': 1.0, 'H+': 1e-7, 'OH-': 1e-7})\n",
    "res = odesys.integrate(tout, c0, atol=1e-12, rtol=1e-14, integrator='gsl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_result(result):\n",
    "    fig, axes = plt.subplots(1, 4, figsize=(16, 6))\n",
    "    result.plot(names=[k for k in rsys.substances if k != 'H2O'], ax=axes[0])\n",
    "    axes[0].legend(loc='best', prop={'size': 10}); _ = axes[0].set_xlabel('Time'); _ = axes[0].set_ylabel('Concentration')\n",
    "    result.plot(names=[k for k in rsys.substances if k != 'H2O'], ax=axes[1], xscale='log', yscale='log')\n",
    "    axes[1].legend(loc='best', prop={'size': 10}); _ = axes[1].set_xlabel('Time'); _ = axes[1].set_ylabel('Concentration')\n",
    "    result.plot(deriv=True, names=[k for k in rsys.substances if k != 'H2O'], ax=axes[2], xscale='log'); \n",
    "    axes[2].set_yscale('symlog', linthresh=1e-9)\n",
    "    axes[2].legend(loc='best', prop={'size': 10}); _ = axes[2].set_xlabel('Time'); _ = axes[2].set_ylabel('dC/dt')\n",
    "    result.plot_invariant_violations(ax=axes[3], xscale='log')\n",
    "    axes[3].legend(loc='best', prop={'size': 10}); _ = axes[3].set_xlabel('Time'); _ = axes[3].set_ylabel('Component conservation')\n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_result(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "substance_keys = 'Fe+2 Fe+3 FeOOH(s) H2O2'.split()\n",
    "selection = res.xout > 10\n",
    "fig, axes = plt.subplots(len(substance_keys), 2, figsize=(16, 4*len(substance_keys)))\n",
    "plot_reaction_contributions(res, rsys, extra['rate_exprs_cb'], substance_keys, axes=axes[:, 0])\n",
    "plot_reaction_contributions(res, rsys, extra['rate_exprs_cb'],\n",
    "                            substance_keys, axes=axes[:, 1], relative=True, xscale='linear', yscale='linear',\n",
    "                            combine_equilibria=True, selection=selection)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nativesys = get_native(rsys, odesys, 'cvode')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nativeres = nativesys.integrate(tout[-1], c0, atol=1e-10, rtol=1e-10, nsteps=25000,\n",
    "                                return_on_root=True, first_step=1e-40) #error_outside_bounds=True)\n",
    "plot_result(nativeres)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from _ipynb_progressbar import log_progress\n",
    "native = get_native(rsys, odesys, 'cvode', steady_state_root=True)\n",
    "def plot_tss_conv(factor, tols, ax, idx):\n",
    "    success = []\n",
    "    failure = []\n",
    "    c = 'krbgcmy'\n",
    "    for tol in log_progress(tols):\n",
    "        res = native.integrate(\n",
    "            1e11, c0, atol=tol, rtol=tol, nsteps=25000,\n",
    "            return_on_root=True, dx_min=1e-17, return_on_error=True, first_step=1e-40, special_settings=[factor])\n",
    "        if res.info['success']:\n",
    "            success.append((tol, res.xout[-1]))\n",
    "        else:\n",
    "            failure.append((tol, res.xout[-1]))\n",
    "    if success:\n",
    "        ax.loglog(*zip(*success), marker='.', color=c[idx % len(c)], label=factor)\n",
    "    if failure:\n",
    "        ax.loglog(*zip(*failure), marker='x', color=c[idx % len(c)], label=None if success else fac01tor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#fig, ax = plt.subplots(figsize=(14, 6))\n",
    "#tols = np.logspace(-10, -6, 50)\n",
    "#for idx, factor in enumerate([1e4, 1.1e4, 1e5, 1e6, 1e7, 1e8, 1e9]):\n",
    "#    plot_tss_conv(factor, tols, ax, idx)\n",
    "#ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dominant_reactions_graph(res.yout[-1, :], extra['rate_exprs_cb'], rsys, 'H2O2', combine_equilibria=True,\n",
    "                         fname='H2O2.png', output_dir='.', include_inactive=False, tex=False)\n",
    "Image('H2O2.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
